#!/usr/bin/env python2.7
#-*- coding:utf-8 -*-
#Author='He Rensheng'
#Email='hrs323@126.com'

import os
from pascal_voc_io import PascalVocWriter
import cv2


def generate_background_xml(imagesdir, image_name, savedir):
    im_file = os.path.join(imagesdir, image_name + '.jpg')
    im = cv2.imread(im_file)
    xml = PascalVocWriter('VOC2007', image_name + '.jpg', im.shape)
    # 将整张图片标记为'__background__'
    xml.addBndBox(1, 1, int(im.shape[1]) - 1, int(im.shape[0]) - 1, '__background__')
    # 保存xml
    xml_file = os.path.join(savedir, image_name + '.xml')
    xml.save(xml_file)

def get_file_list(imagesdir):
    fileList = []
    for file in os.listdir(imagesdir):
        if file[-4:] == '.jpg':
            fileList.append(file[:-4])
    return fileList


if __name__ == '__main__':
    imagesdir = '/home/hrs/installations/pva-faster-rcnn/data/VOCdevkit2007/VOC2007/JPEGImages'
    savedir = '/home/hrs/installations/pva-faster-rcnn/data/VOCdevkit2007/VOC2007/background_xml/'
    if not os.path.exists(savedir):
        os.makedirs(savedir)

    imgfilelist = get_file_list(imagesdir)
    for imgfile in imgfilelist:
        generate_background_xml(imagesdir, imgfile, savedir)